# Bart模型 pip install textgen
from transformers import BertTokenizerFast
from textgen import BartSeq2SeqModel
tokenizer = BertTokenizerFast.from_pretrained('shibing624/bart4csc-base-chinese')
m = BartSeq2SeqModel(
    encoder_type='bart',
    encoder_decoder_type='bart',
    encoder_decoder_name='shibing624/bart4csc-base-chinese',
    tokenizer=tokenizer,
    args={"max_length": 128, "eval_batch_size": 128})

def demo2():

    tt = '一级域名数：1个，子域名数量：10个，黑白名单个数：20个\n支持HTTP、HTTPS的业务防护（不限于80、8080、443、8443端口），支持针对攻击频繁的恶意IP自动封进，支持防护如SQL注入、命令执行等常见WEB攻击'
    idx_errors = m.predict([tt])
    print(idx_errors)

if __name__ == '__main__':
    demo2()